#pragma once
#include "op_common.h"
#include <cstdint>
namespace AscendC {
// 参考：tensorflow/core/framework/ops_util.h
// 搜索：ComputeStride
// Helper to compute 'strides' given a tensor 'shape'. I.e.,
// strides[i] = prod(shape.dim_size[(i+1):])
template <typename T>
__aicore__ inline void ComputeStride(T *shape, int64_t ndims, T *strides) {
  T stride = 1;
  for (int i = ndims - 1; i >= 0; --i) {
    strides[i] = stride;
    stride *= shape[i];
  }
}

// 参考：tensorflow/core/kernels/transpose_functor_cpu.cc
// 搜索: TransposeSimple
template <typename T>
__aicore__ inline void
TransposeSimple(__gm__ T *in, int64_t size, int64_t *in_shape, int64_t ndims,
                int64_t *out_shape, int64_t *perm, __gm__ T *out) {
  int64_t in_strides[MAX_SHAPE_DIM] = {0};
  int64_t out_strides[MAX_SHAPE_DIM] = {0};
  ComputeStride(in_shape, ndims, in_strides);
  ComputeStride(out_shape, ndims, out_strides);
  for (int64_t o_idx = 0; o_idx < size; ++o_idx) {
    int64_t i_idx = 0;
    int64_t t = o_idx;
    for (int i = 0; i < ndims; ++i) {
      const int64_t ratio = t / out_strides[i];
      t -= ratio * out_strides[i];
      i_idx += ratio * in_strides[perm[i]];
    }
    out[o_idx] = in[i_idx];
  }
}

struct PermuteHelper {
  int64_t ndims_ = 0;
  int64_t *perm_ = nullptr;
  int64_t from_strides_[MAX_SHAPE_DIM] = {0};
  int64_t to_strides_[MAX_SHAPE_DIM] = {0};
  __aicore__ inline PermuteHelper() {}
  __aicore__ inline void init(int64_t *fromShape, int64_t ndims,
                              int64_t *toShape, int64_t *perm) {
    ndims_ = ndims;
    perm_ = perm;
    ComputeStride(fromShape, ndims, from_strides_);
    ComputeStride(toShape, ndims, to_strides_);
  }
  __aicore__ inline int64_t get(int64_t o_idx) {
    int64_t i_idx = 0;
    int64_t t = o_idx;
    for (int i = 0; i < ndims_; ++i) {
      const int64_t ratio = t / to_strides_[i];
      t -= ratio * to_strides_[i];
      i_idx += ratio * from_strides_[perm_[i]];
    }
    return i_idx;
  }
};
} // namespace AscendC